/**
 * @file aimg.c
 * @author ailen (noorim.ailen@gmail.com)
 * @version 0.2.3
 * @modified 10.02.17
 * @brief uncomment 'define KNL_ON' to contain your kernel
 */
#include <stdio.h>
#include <assert.h>

/**
 * some macros
 */
#define FLP_SIZE	(sizeof(BYTE) * 1440 * 1024)
#define KNL_ON
#ifdef _cplusplus
#define INLINE inline
#else
#define INLINE
#endif
/**
 * type definitions
 */
typedef unsigned char	BYTE;
typedef unsigned int	UINT;
typedef unsigned short	USHORT;

typedef struct ImgSliceTag
{
	BYTE	*pBuff;
	UINT	iSize;
	USHORT	nSect;
	UINT	iPos;
} ImgSlice;

typedef struct ImageTag
{
	ImgSlice	isBoot;
	ImgSlice	isKload;
	ImgSlice	isKnl;
} Image;

int fsize(FILE *pFileHandle);
int loadSlice(ImgSlice *pImgSlice, char *filename);
int loadImage(Image *pImg, char *boot_filename, char *kload_filename, char *kernel_filename);
int writeImage(Image *pImg, char *img_filename);
void error(void);

///==================================================================

/**
 * @param argument's count
 * @param argument's values
 * @brief entrance
 */
int main(int argc, char *argv[])
{
	Image img;
	//printf("AImg v0.2.2\n"); /// fixed at 10.7.18

	if (argc != 4)
	{
		fputs("Usage: aimg.exe <boot loader> <kload> <kernel>\n", stderr); /// fixed at 10.7.18
		error();
	}

	if ( loadImage(&img, argv[1], argv[2], argv[3]) != 0 ) error();
	if ( writeImage(&img, "ailenos.img") != 0 ) error();
	if ( cleanImage(&img) != 0 ) error();

	return 0;
}

/**
 * @param pImg thing to load
 * @param boot_filename
 * @param kload_filename
 * @param kernel_filename
 */
int loadImage(Image *pImg, char *boot_filename, char *kload_filename, char *kernel_filename)
{
#ifdef _DEBUG
	printf("Name\t\tSize\tSector\n");
#endif
	if (loadSlice(&pImg->isBoot, boot_filename) != 0) return -1;
	if (loadSlice(&pImg->isKload, kload_filename) != 0) goto failed_kload;
#ifdef KNL_ON
	if (loadSlice(&pImg->isKnl, kernel_filename) != 0) goto failed_kernel;
#endif
	return 0;

	/// exception processing
failed_kernel:	cleanSlice(&pImg->isKload);
failed_kload:	cleanSlice(&pImg->isBoot);
	return -1;
}

/**
 * @param pImg
 * @param img_filename
 */
int writeImage(Image *pImg, char *img_filename)
{
	USHORT	sSect;
	BYTE	*pMemory;
	FILE	*fImg;
#ifdef KNL_ON
	sSect = pImg->isKload.nSect + pImg->isKnl.nSect;
#else
	sSect = pImg->isKload.nSect;
#endif
	//printf("Whole Sectors: %d\n", sSect);

	pImg->isBoot.pBuff[5] = sSect & 0xff;
	pImg->isBoot.pBuff[6] = sSect >> 8; // Bug Point 1st

	pMemory = malloc(FLP_SIZE);
	if (pMemory == NULL)
	{
		printf("Error: Cannot allocate memory.");
		cleanImage(pImg);
		return -1;
	}
	pImg->isBoot.iPos = 0;
	pImg->isKload.iPos = pImg->isBoot.iPos + pImg->isBoot.nSect * 512;
	pImg->isKnl.iPos = pImg->isKload.iPos + pImg->isKload.nSect * 512;

	memmove(&pMemory[pImg->isBoot	.iPos], pImg->isBoot	.pBuff, pImg->isBoot	.iSize);
	memmove(&pMemory[pImg->isKload	.iPos], pImg->isKload	.pBuff, pImg->isKload	.iSize);
#ifdef KNL_ON
	memmove(&pMemory[pImg->isKnl	.iPos], pImg->isKnl		.pBuff, pImg->isKload	.iSize);
#endif
	if ( (fImg = fopen("ailenos.img", "wb")) == NULL )
	{
		printf("Error: Cannot open the file:ailenos.img\n");
		free(pMemory);
		return -1;
	}

	fwrite(pMemory, FLP_SIZE, 1, fImg);
	fclose(fImg);
	free(pMemory);
	return 0;
}

/**
 * @brief clean image
 */
int cleanImage(Image *pImg)
{
	cleanSlice(&pImg->isBoot);
	cleanSlice(&pImg->isKload);
#ifdef KNL_ON
	cleanSlice(&pImg->isKnl);
#endif
	return 0;
}

/**
 * @param pImgSlice to load
 * @param filename to read
 */
int loadSlice(ImgSlice *pImgSlice, char *filename)
{
	FILE	*pFile;
	if ( (pFile = fopen(filename, "rb")) == NULL )
	{
		printf("Error: Cannot open the file:%s\n", filename);
		return -1;
	}

	pImgSlice->iSize = fsize(pFile);
	pImgSlice->nSect = ( pImgSlice->iSize + 511 )/512;
	pImgSlice->pBuff = malloc(pImgSlice->iSize);
	if ( pImgSlice->pBuff == NULL )
	{
		printf("Error: Cannot allocate memory.");
		fclose(pFile);
		return -1;
	}
	
	fread(pImgSlice->pBuff, 512, pImgSlice->nSect, pFile);
	fclose(pFile);

#ifdef _DEBUG
	fprintf(stderr, "%s\t%7d\t%7d\n", filename, pImgSlice->iSize, pImgSlice->nSect);
#endif
	return 0;
}

INLINE int cleanSlice(ImgSlice *pImgSlice)
{
	free(pImgSlice->pBuff);
	return 0;
}

/**
 * @brief google as `file size in c'
 * @TODO support 64bit
 */
int fsize(FILE *pFileHandle)
{
	int iR;

	fseek(pFileHandle, 0, SEEK_END);
	iR = ftell(pFileHandle);
	fseek(pFileHandle, 0, SEEK_SET);
	return iR;
}

/**
 * @brief puts error end exit
 */
INLINE void error(void)
{
	printf("error!\n");
	exit(-1);
}
